import numpy as np
import matplotlib.pyplot as plt

# ---------- Coordinate Conversion ----------

def ra_dec_to_mollweide_x_y(ra_deg, dec_deg):
        """
        Convert (RA, Dec) to Mollweide projection coordinates (x,y)
        - RA input in degrees [0, 360), astronomical convention: RA increases to the left
        - Returns (x,y) in radians, range x∈[-pi,pi], y∈[-pi/2,pi/2]
        """
        ra = np.radians(np.asarray(ra_deg))
        dec = np.radians(np.asarray(dec_deg))
        # Flip RA direction: x = -(RA - pi)
        x = -(ra - np.pi)
        # Wrap to [-pi, pi]
        x = (x + np.pi) % (2*np.pi) - np.pi
        return x, dec

def setup_ra_hour_xticks(ax):
    """
    Beautify x-axis ticks: use RA hours (increasing from right to left)
    """
    tick_deg = np.arange(-150, 180, 30)         # Mollweide longitude ticks
    tick_rad = np.radians(tick_deg)
    ax.set_xticks(tick_rad)
    labels = []
    for d in tick_deg:
        # RA(deg) = (180 - longitude), then /15 to get hours
        ra_deg_here = (180 - d) % 360
        h = int(round(ra_deg_here / 15.0)) % 24
        labels.append(f"{h:02d}h")
    ax.set_xticklabels(labels)

# ---------- Draw Ecliptic ----------

def plot_ecliptic(ax, color="r", lw=1.5, label="Ecliptic"):
    """
    Plot ecliptic (J2000) on existing mollweide axes using only numpy
    dec(ra) ≈ asin(sin(eps)*sin(ra)), eps=23.44°
    """
    eps = np.radians(23.44)  # Obliquity of the ecliptic
    ra = np.linspace(0, 2*np.pi, 2000)         # RA in rad
    dec = np.arcsin(np.sin(eps) * np.sin(ra))  # Dec in rad
    # Convert to Mollweide x-coordinate and flip direction
    x = -(ra - np.pi)
    x = (x + np.pi) % (2*np.pi) - np.pi
    ax.plot(x, dec, color=color, lw=lw, label=label)

# ---------- Main Plotting Function ----------

def plot_sky_coverage_with_ecliptic(
    ra_deg_list,
    dec_deg_list,
    title="Sky Coverage (J2000) with Ecliptic",
    save_path=None,
    green_box=None,   # ((ra_c_deg, dec_c_deg), width_arcmin, height_arcmin)
    mark_point=None   # (ra_deg, dec_deg, text)
):
    fig = plt.figure(figsize=(12, 6))
    ax = fig.add_subplot(111, projection="mollweide")

    # Sky coverage points
    x, y = ra_dec_to_mollweide_x_y(ra_deg_list, dec_deg_list)
    ax.scatter(x, y, s=4, alpha=0.8)

    # Overlay ecliptic
    plot_ecliptic(ax, color="tab:red", lw=1.6, label="Ecliptic")

    # Optional: Draw green box (given center and width/height in arcmin)
    if green_box is not None:
        (ra_c, dec_c), w_arcmin, h_arcmin = green_box
        # Sample rectangle edges, convert to projection and plot
        w_deg = w_arcmin / 60.0
        h_deg = h_arcmin / 60.0
        # Approximate rectangle with polyline
        # Top edge
        edge_ra = np.linspace(ra_c - w_deg/2, ra_c + w_deg/2, 200)
        edge_dec = np.full_like(edge_ra, dec_c + h_deg/2)
        # Bottom edge
        edge_ra2 = np.linspace(ra_c + w_deg/2, ra_c - w_deg/2, 200)
        edge_dec2 = np.full_like(edge_ra2, dec_c - h_deg/2)
        # Left and right edges
        edge_ra3 = np.full(200, ra_c - w_deg/2)
        edge_dec3 = np.linspace(dec_c - h_deg/2, dec_c + h_deg/2, 200)
        edge_ra4 = np.full(200, ra_c + w_deg/2)
        edge_dec4 = np.linspace(dec_c + h_deg/2, dec_c - h_deg/2, 200)

        for er, ed in [(edge_ra, edge_dec), (edge_ra2, edge_dec2),
                       (edge_ra3, edge_dec3), (edge_ra4, edge_dec4)]:
            bx, by = ra_dec_to_mollweide_x_y(er, ed)
            ax.plot(bx, by, color="limegreen", lw=1.2)

    # Optional: Mark a point
    if mark_point is not None:
        pra, pdec, ptxt = mark_point
        px, py = ra_dec_to_mollweide_x_y(pra, pdec)
        ax.scatter(px, py, s=28, color="tab:blue", zorder=5)
        ax.text(px, py, " "+ptxt, fontsize=9, va="center")

    # Coordinate system decoration
    setup_ra_hour_xticks(ax)
    yticks_deg = [-75, -45, -15, 15, 45, 75]
    ax.set_yticks(np.radians(yticks_deg))
    ax.set_yticklabels([f"{d}°" for d in yticks_deg])

    ax.grid(True, alpha=0.5)
    ax.legend(loc="lower left", framealpha=0.7)
    plt.title(title)

    if save_path:
        plt.savefig(save_path, dpi=200, bbox_inches="tight")
    plt.show()

# ---------- 示例调用 ----------
if __name__ == "__main__":
    # 这里用一些随机点来模拟你的覆盖；实际使用时替换成你的 RA/Dec 数组（单位：度）
    rng = np.random.default_rng(42)
    ra_demo = (rng.random(500)*360)                 # [0,360)
    dec_demo = (rng.random(500)*180 - 90)           # [-90,90]

    # 想画绿色框：以 06:10:38 +21:38:42 为中心，宽 60′ 高 60′
    ra_point_deg = 15*(6 + 10/60 + 38/3600)         # 06:10:38 -> deg
    dec_point_deg = 21 + 38/60 + 42/3600

    plot_sky_coverage_with_ecliptic(
        ra_deg_list=ra_demo,
        dec_deg_list=dec_demo,
        title="Sky Coverage + Ecliptic (Mollweide, J2000)",
        save_path="/mnt/7b21f1e1-eb25-4cd5-bdb5-06d7d82fa253/Temp/force_photmetry/plot_py/demo02.png",  # "sky_with_ecliptic.png"
        green_box=((ra_point_deg, dec_point_deg), 60, 60),
        mark_point=(ra_point_deg, dec_point_deg, "06:10:38  +21°38′42″")
    )
